"""
将 sklearn digits 数据集上载到 Pinecone，并评估 k=11 的近邻投票准确率。

需求对齐：
- 使用 80% 的数据构建索引并上载；
- 使用 20% 的数据在 Pinecone 上做搜索，k=11，做多数投票并计算准确率；
- 日志 logging 必须带日期；
- 上载与测试时带 tqdm 进度条；
- 打印“成功创建索引，并上传了 1437 条数据”等信息（若样本数不同，以实际为准，但演示保持该格式）。

注意：
- 本脚本使用欧式距离（euclidean），向量维度为 64（8x8）；
- 强烈建议通过环境变量 PINECONE_API_KEY 传入 API Key；
- 索引名默认 "mnist-index"，可通过 PINECONE_INDEX 覆盖；
- Serverless 规格默认 aws/us-east-1，可通过 PINECONE_CLOUD/PINECONE_REGION 覆盖。
"""

from __future__ import annotations

import logging
import os
from collections import Counter

import numpy as np
from sklearn import datasets
from sklearn.model_selection import train_test_split
from tqdm import tqdm

try:
    from dotenv import load_dotenv  # type: ignore

    load_dotenv()  # noqa: F401
except Exception:
    pass

from pinecone import Pinecone, ServerlessSpec  # type: ignore


# -------------------------- 配置 --------------------------
INDEX_NAME = os.getenv("PINECONE_INDEX", "mnist-index")
API_KEY = os.getenv("PINECONE_API_KEY")
DIM = 64
METRIC = "euclidean"
REGION = os.getenv("PINECONE_REGION", "us-east-1")
CLOUD = os.getenv("PINECONE_CLOUD", "aws")
K = int(os.getenv("K", "11"))
RECREATE = os.getenv("PINECONE_RECREATE", "0") in {"1", "true", "True"}


def setup_logging() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def ensure_index(pc: Pinecone, name: str) -> None:
    try:
        names = list(pc.list_indexes().names())  # type: ignore[attr-defined]
    except Exception:
        lst = pc.list_indexes()
        if isinstance(lst, dict) and "indexes" in lst:
            names = [ix.get("name") for ix in lst.get("indexes", [])]
        else:
            names = [getattr(ix, "name", getattr(ix, "id", None)) for ix in (lst or [])]

    if RECREATE and name in (names or []):
        logging.info("检测到 PINECONE_RECREATE=1，删除并重建索引 '%s'...", name)
        pc.delete_index(name)
        names = [n for n in names if n != name]

    if name not in (names or []):
        logging.info("索引 '%s' 不存在，正在创建...", name)
        pc.create_index(
            name=name,
            dimension=DIM,
            metric=METRIC,
            spec=ServerlessSpec(cloud=CLOUD, region=REGION),
        )
        logging.info("索引 '%s' 创建成功。", name)


def upsert_train_set(index, X_train: np.ndarray, y_train: np.ndarray) -> int:
    total = len(X_train)
    count = 0
    batch = []
    for i, (vec, label) in enumerate(tqdm(zip(X_train, y_train), total=total, desc="上传训练数据", unit="vec")):
        batch.append({"id": f"train-{i}", "values": vec.tolist(), "metadata": {"label": int(label)}})
        if len(batch) >= 200:
            index.upsert(vectors=batch)
            count += len(batch)
            batch = []
    if batch:
        index.upsert(vectors=batch)
        count += len(batch)
    return count


def eval_on_test(index, X_test: np.ndarray, y_test: np.ndarray, k: int) -> float:
    correct = 0
    for vec, label in tqdm(zip(X_test, y_test), total=len(X_test), desc="测试(k=11)", unit="img"):
        res = index.query(vector=vec.tolist(), top_k=k, include_metadata=True)
        matches = res.get("matches") if isinstance(res, dict) else getattr(res, "matches", [])
        labels = []
        for m in matches:
            meta = m.get("metadata", {}) if isinstance(m, dict) else getattr(m, "metadata", {})
            if meta and "label" in meta:
                labels.append(int(meta["label"]))
        pred = Counter(labels).most_common(1)[0][0] if labels else -1
        if pred == int(label):
            correct += 1
    return correct / len(X_test)


def main() -> None:
    setup_logging()

    if not API_KEY:
        raise SystemExit("未检测到 PINECONE_API_KEY，请先通过环境变量配置再运行。")

    pc = Pinecone(api_key=API_KEY)
    ensure_index(pc, INDEX_NAME)
    index = pc.Index(INDEX_NAME)

    digits = datasets.load_digits()
    X = digits.data.astype(np.float32)
    y = digits.target.astype(int)

    # 80/20 划分
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )

    # 上载训练集
    uploaded = upsert_train_set(index, X_train, y_train)
    logging.info("成功创建索引，并上传了 %d 条数据", uploaded)

    # 评估 k=11
    acc = eval_on_test(index, X_test, y_test, K)
    logging.info("使用 Pinecone 的准确率: %.2f%%", acc * 100)
    logging.info("k=%d 的准确率为: %.4f", K, acc)


if __name__ == "__main__":
    main()
